import torch.nn as nn
from bert_config_zc import Bert_Config_ZC
from transformers import BertModel
import torch.nn.functional as F

conf = Bert_Config_ZC()

class PediatricClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.bertmodel = BertModel.from_pretrained(conf.bert_pretrain_path)
        self.fc = nn.Linear(conf.bert_conf.hidden_size, 256)
        self.out = nn.Linear(256, conf.num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bertmodel(input_ids=input_ids, attention_mask=attention_mask)
        pooler_output = outputs['pooler_output']
        logits = self.out(F.relu(self.fc(pooler_output)))
        # logits = self.fc(outputs['pooler_output'])
        return logits


